const std = @import("std");
const builtin = @import("builtin");
const mem = std.mem;
const Ast = std.zig.Ast;
const Node = Ast.Node;
const Access = enum {
    none,
    public,
    private,
    c,
};
const ScopeList = std.ArrayList(struct {
    kind: []const u8,
    scope: []const u8,
});
const Error = std.Io.Writer.Error || error{OutOfMemory};

var ast: Ast = undefined;
var stack: ScopeList = .empty;
var filename: []const u8 = undefined;

var writer: *std.Io.Writer = undefined;
var debug_allocator: std.heap.DebugAllocator(.{}) = .init;
var gpa: mem.Allocator = undefined;

const Signature = struct {
    parm_list: []const u8,
    return_type: ?[]const u8 = null,
};

pub fn main() anyerror!void {
    gpa, const is_debug = gpa: {
        if (builtin.os.tag == .wasi) break :gpa .{ std.heap.wasm_allocator, false };
        break :gpa switch (builtin.mode) {
            .Debug, .ReleaseSafe => .{ debug_allocator.allocator(), true },
            .ReleaseFast, .ReleaseSmall => .{ std.heap.smp_allocator, false },
        };
    };
    defer if (is_debug) {
        _ = debug_allocator.deinit();
    };

    const args = try std.process.argsAlloc(gpa);
    defer std.process.argsFree(gpa, args);

    if (args.len < 2) {
        std.debug.print("Usage: {s} file.zig\n", .{args[0]});
        return;
    }

    filename = args[1];
    const source = try std.fs.cwd().readFileAllocOptions(
        gpa,
        filename,
        std.math.maxInt(usize),
        null,
        .@"1",
        0,
    );
    defer gpa.free(source);

    const mode: Ast.Mode = if (std.mem.endsWith(u8, filename, ".zig")) .zig else .zon;
    ast = try .parse(gpa, source, mode);
    defer ast.deinit(gpa);
    defer stack.deinit(gpa);

    var stdout_buffer: [4096]u8 = undefined;
    var stdout_writer = std.fs.File.stdout().writer(&stdout_buffer);
    writer = &stdout_writer.interface;

    // root ContainerMembers
    for (ast.rootDecls()) |member| {
        try printTags(member);
    }

    try writer.flush();
}

fn printTags(index: Node.Index) !void {
    const main_token = ast.nodeMainToken(index);
    const data = ast.nodeData(index);

    switch (ast.nodeTag(index)) {
        .global_var_decl,
        .local_var_decl,
        .simple_var_decl,
        .aligned_var_decl,
        => var_decl: {
            const full = ast.fullVarDecl(index).?;
            const access: Access = if (full.extern_export_token) |_|
                .c
            else if (full.visib_token) |_|
                .public
            else
                .none;

            if (full.ast.init_node.unwrap()) |init_node| {
                var buf: [2]Ast.Node.Index = undefined;
                const init_node_decl = ast.fullContainerDecl(&buf, init_node);
                if (init_node_decl) |decl| {
                    try printContainer(main_token + 1, decl, access, null);
                    break :var_decl;
                }

                switch (ast.nodeTag(init_node)) {
                    .error_set_decl,
                    .merge_error_sets,
                    => {
                        try printLine(.{
                            .tag = main_token + 1,
                            .kind = "error",
                            .access = access,
                        });
                        break :var_decl;
                    },

                    .builtin_call_two => {
                        const init_node_token = ast.tokenSlice(ast.nodeMainToken(init_node));
                        if (mem.eql(u8, init_node_token, "@import") or mem.eql(u8, init_node_token, "@cImport")) {
                            try printLine(.{
                                .tag = main_token + 1,
                                .kind = "import",
                                .access = access,
                            });
                            break :var_decl;
                        }
                    },

                    else => {},
                }
            }

            // `var` or `const`
            try printLine(.{
                .tag = main_token + 1,
                .kind = ast.tokenSlice(main_token),
                .access = access,
            });
        },

        .fn_proto_simple,
        .fn_proto_multi,
        .fn_proto_one,
        .fn_proto,
        .fn_decl,
        => |tag| fn_decl: {
            var buffer: [1]Node.Index = undefined;
            const full = ast.fullFnProto(&buffer, index).?;
            const public = if (full.visib_token) |_| true else false;
            const access: Access = if (full.extern_export_inline_token) |t|
                if (std.mem.eql(u8, ast.tokenSlice(t), "inline"))
                    if (public) .public else .private
                else
                    .c
            else if (public) .public else .none;

            // get function signature
            const token_tags = ast.tokens.items(.tag);
            const l_paren = full.lparen;
            const r_paren = rp: {
                var n: isize = 0;

                break :rp for (token_tags[l_paren + 1 ..], l_paren + 1..) |token, i| {
                    switch (token) {
                        .l_paren => n += 1,
                        .r_paren => {
                            n -= 1;
                            if (n < 0) break i;
                        },
                        else => {},
                    }
                } else unreachable;
            };

            const token_starts = ast.tokens.items(.start);
            const l_start = token_starts[l_paren];
            const r_start = token_starts[r_paren];
            const parm_list = ast.source[l_start + 1 .. r_start];

            const last_token = ast.lastToken(full.ast.proto_node) + 1;
            const return_type = ast.source[r_start + 1 .. token_starts[last_token]];

            // HACK: a function may return a struct, we treat it as a container
            if (tag == .fn_decl and std.mem.eql(u8, std.mem.trim(u8, return_type, " "), "type")) {
                var buf: [2]Ast.Node.Index = undefined;
                const container_desc: ?Ast.full.ContainerDecl = cd: {
                    const fn_body_node = data.node_and_node[1];
                    const statements = ast.blockStatements(&buf, fn_body_node) orelse break :cd null;
                    const last_statement = statements[statements.len - 1];

                    if (ast.nodeTag(last_statement) == .@"return") {
                        if (ast.nodeData(last_statement).opt_node.unwrap()) |return_value| {
                            break :cd ast.fullContainerDecl(&buf, return_value);
                        }
                    }
                    break :cd null;
                };

                if (container_desc) |desc| {
                    try printContainer(main_token + 1, desc, access, parm_list);
                    break :fn_decl;
                }
            }

            try printLine(.{
                .tag = main_token + 1,
                .kind = "function",
                .access = access,
                .signature = .{
                    .parm_list = parm_list,
                    .return_type = return_type,
                },
            });
        },

        .container_field_init,
        .container_field_align,
        .container_field,
        => try printLine(.{
            .tag = main_token,
            .kind = "field",
        }),

        .test_decl => try printLine(.{
            .tag = data.opt_token_and_node[0].unwrap() orelse main_token,
            .kind = "test",
        }),

        .@"comptime" => try printLine(.{
            .tag = main_token,
            .kind = "comptime",
        }),

        .struct_init_one,
        .struct_init_one_comma,
        .struct_init_dot_two,
        .struct_init_dot_two_comma,
        .struct_init_dot,
        .struct_init_dot_comma,
        .struct_init,
        .struct_init_comma,
        => {
            // print zon fields
            try printStructFields(index);
        },

        else => |unknown_tag| std.log.debug(
            "unknown: \x1b[33m{s}\x1b[m",
            .{@tagName(unknown_tag)},
        ),
    }
}

fn printContainer(
    tag: Ast.TokenIndex,
    container: Ast.full.ContainerDecl,
    access: Access,
    parm_list: ?[]const u8,
) Error!void {
    const kind = ast.tokenSlice(container.ast.main_token);

    // const A = `struct {}`, `union {}`, `enum {}` or `opaque {}`
    try printLine(.{
        .tag = tag,
        .kind = kind,
        .access = access,
        .signature = if (parm_list == null) null else .{ .parm_list = parm_list.? },
    });

    try stack.append(gpa, .{
        .kind = kind,
        .scope = ast.tokenSlice(tag),
    });
    defer _ = stack.pop();

    // print tags of their ContainerMembers
    for (container.ast.members) |member| {
        try printTags(member);
    }
}

fn printStructFields(index: Node.Index) !void {
    const first_token = ast.firstToken(index);
    const field_token = if (first_token > 0) first_token - 2 else null;

    var buf: [2]Ast.Node.Index = undefined;
    const struct_decl = ast.fullStructInit(&buf, index) orelse {
        try printLine(.{
            .tag = field_token.?,
            .kind = "field",
        });
        return;
    };

    if (field_token) |token| {
        try printLine(.{
            .tag = token,
            .kind = "struct",
        });

        try stack.append(gpa, .{
            .kind = "struct",
            .scope = ast.tokenSlice(token),
        });
    }

    defer if (field_token != null) {
        _ = stack.pop();
    };

    for (struct_decl.ast.fields) |field| {
        try printStructFields(field);
    }
}

fn printLine(info: struct {
    tag: Ast.TokenIndex,
    kind: []const u8,
    access: Access = .none,
    signature: ?Signature = null,
}) !void {
    const loc = ast.tokenLocation(0, info.tag);

    if (mem.eql(u8, info.kind, "comptime")) {
        try writer.print("comptime_{}", .{loc.line + 1});
    } else {
        try writer.print("{s}", .{ast.tokenSlice(info.tag)});
    }

    try writer.print(
        "\t{[file]s}\t{[line]};\"\t{[kind]s}\tline:{[line]}\tcolumn:{[column]}",
        .{
            .file = filename,
            .kind = info.kind,
            .line = loc.line + 1,
            .column = loc.column + 1,
        },
    );

    // write scopes
    if (stack.items.len > 0) {
        try writer.print("\t{s}:", .{stack.getLast().kind});

        try writer.print("{s}", .{stack.items[0].scope});
        for (stack.items[1..]) |scope| {
            try writer.print(".{s}", .{scope.scope});
        }
    }

    switch (info.access) {
        .public, .private => |tag| try writer.print("\taccess:{s}", .{@tagName(tag)}),
        // HACK: Treat C ABI functions as 'protected' for tagbar visualization.
        // Zig doesn't have protected visibility, this is just for UI purposes.
        .c => try writer.writeAll("\taccess:protected"),
        else => {},
    }

    if (info.signature) |s| {
        const return_type_len = if (s.return_type) |rt| rt.len else 0;
        var buffer: std.ArrayList(u8) = try .initCapacity(gpa, @max(s.parm_list.len, return_type_len));
        defer buffer.deinit(gpa);

        trim(&buffer, s.parm_list);
        try writer.print("\tsignature: ({s})", .{mem.trimRight(u8, buffer.items, ", ")});

        if (s.return_type) |rt| {
            buffer.clearRetainingCapacity();
            trim(&buffer, rt);
            try writer.print(" {s}", .{mem.trimRight(u8, buffer.items, " ")});
        }
    }

    try writer.writeByte('\n');
}

// trim newline and comments
fn trim(w: *std.ArrayList(u8), string: []const u8) void {
    var it = mem.tokenizeScalar(u8, string, '\n');
    while (it.next()) |line| {
        const end = mem.indexOf(u8, line, "//") orelse line.len;
        const str = mem.trim(u8, line[0..end], &std.ascii.whitespace);
        if (str.len == 0) continue;

        w.print(gpa, "{s} ", .{str}) catch unreachable;
    }
}
